#include <TRRT_ompl/TRRT.h>
#include <ompl/base/GoalSampleableRegion.h>
#include <ompl/datastructures/NearestNeighborsSqrtApprox.h>
#include <ompl/tools/config/SelfConfig.h>


namespace ompl
{
namespace geometric
{

void TRRT::clear()
{
  Planner::clear();
  sampler_.reset();
  freeMemory();
  if(nn_)
    nn_->clear();
}

void TRRT::setup()
{
  Planner::setup();
  SelfConfig sc(si_, getName());
  sc.configurePlannerRange(maxDistance_);

  // need to set parameters

  if(!nn_)
    nn_.reset(new NearestNeighborsSqrtApprox<Motion*>());
  nn_->setDistanceFunction(boost::bind(&TRRT::distanceFunction, this, _1, _2));
}

void TRRT::freeMemory()
{
  if(nn_)
  {
    std::vector<Motion*> motions;
    nn_->list(motions);
    for(unsigned int i = 0; i < motions.size(); ++i)
    {
      if(motions[i]->state)
        si_->freeState(motions[i]->state);
      delete motions[i];
    }
  }
}

void TRRT::getPlannerData(base::PlannerData& data) const
{
  Planner::getPlannerData(data);

  std::vector<Motion*> motions;
  if (nn_)
      nn_->list(motions);

  for (unsigned int i = 0 ; i < motions.size() ; ++i)
      data.recordEdge(motions[i]->parent ? motions[i]->parent->state : NULL, motions[i]->state);
}

bool TRRT::solve(const base::PlannerTerminationCondition& ptc)
{
  checkValidity();
  base::Goal* goal = pdef_->getGoal().get();
  base::GoalSampleableRegion* goal_s = dynamic_cast<base::GoalSampleableRegion*>(goal);

  if(!motionCostEstimator_)
    msg_.error("A motion cost estimator must be provided!");

  unsigned int numFail = 0;

  normalizeK_ = maxDistance_; // different from the original TRRT

  while(const base::State* st = pis_.nextStart())
  {
    Motion* motion = new Motion(si_);
    si_->copyState(motion->state, st);
    motion->cost = 0;
    nn_->add(motion);
  }

  if(nn_->size() == 0)
  {
    msg_.error("There are no valid initial states!");
    return false;
  }

  if(!sampler_)
    sampler_ = si_->allocStateSampler();

  msg_.inform("Starting with %u states", nn_->size());

  Motion *solution  = NULL;
  Motion *approxsol = NULL;
  double  approxdif = std::numeric_limits<double>::infinity();
  Motion *rmotion   = new Motion(si_);
  base::State* rstate = rmotion->state;
  base::State* xstate = si_->allocState();


  while(ptc() == false)
  {
    // sample random state (with goal biasing)
    if (goal_s && rng_.uniform01() < goalBias_ && goal_s->canSample())
        goal_s->sampleGoal(rstate);
    else
        sampler_->sampleUniform(rstate);


    // find closest state in the tree
    Motion* nmotion = nn_->nearest(rmotion);
    base::State* dstate = rstate;

    // find state to add
    double d = si_->distance(nmotion->state, rstate);
    if(d > maxDistance_)
    {
      si_->getStateSpace()->interpolate(nmotion->state, rstate, maxDistance_ / d, xstate);
      dstate = xstate;
      d = maxDistance_;
    }

    if(si_->checkMotion(nmotion->state, dstate))
    {
      bool transitionTestSucceed = false;

      double curCost = motionCostEstimator_->computeCost(nmotion->state, dstate) + nmotion->cost;

      base::State* gstate = si_->allocState();
      goal_s->sampleGoal(gstate);

      if(curCost <= costMax_)
      {
        double curF = curCost + si_->distance(dstate, gstate);
        double preF = nmotion->cost + si_->distance(nmotion->state, gstate);
        if(curF < preF)
        {
          transitionTestSucceed = true;
        }
        else
        {
          double p = exp(-(curF - preF) / (d * normalizeK_ * temperature_));
          double r = rng_.uniform01();
          if(r < p)
          {
            temperature_ /= temperatureFactor_;
            numFail = 0;
            transitionTestSucceed = true;
          }
          else
          {
            if(numFail > numFailMax_)
            {
              temperature_ *= temperatureFactor_;
              numFail = 0;
            }
            else
              numFail += 1;
          }
        }
      }

      si_->freeState(gstate);

      if(transitionTestSucceed)
      {
        Motion* motion = new Motion(si_);
        si_->copyState(motion->state, dstate);
        motion->parent = nmotion;
        motion->cost = curCost;

        nn_->add(motion);
        double dist = 0.0;
        bool solved = goal->isSatisfied(motion->state, &dist);
        if(solved)
        {
          approxdif = dist;
          solution = motion;
          break;
        }

        if(dist < approxdif)
        {
          approxdif = dist;
          approxsol = motion;
        }
      }
    }
  }


  bool approximate = false;
  if(solution == NULL)
  {
    solution = approxsol;
    approximate = true;
  }

  if(solution != NULL)
  {
    // construct the solution path
    std::vector<Motion*> mpath;
    while(solution != NULL)
    {
      mpath.push_back(solution);
      solution = solution->parent;
    }

    // set the solution path
    PathGeometric* path = new PathGeometric(si_);
    for(int i = mpath.size() - 1 ; i >= 0 ; --i)
      path->states.push_back(si_->cloneState(mpath[i]->state));

    goal->setDifference(approxdif);
    goal->setSolutionPath(base::PathPtr(path), approximate);

    if (approximate)
        msg_.warn("Found approximate solution");

    msg_.inform("Cost of the path is %f", mpath[0]->cost);
    double length = 0;
    for(unsigned int i = 0; i < mpath.size() - 1; ++i)
    {
      length += si_->distance(mpath[i]->state, mpath[i + 1]->state);
    }

    msg_.inform("Length of the path is %f", length);
  }


  si_->freeState(xstate);
  if(rmotion->state)
    si_->freeState(rmotion->state);
  delete rmotion;

  msg_.inform("Created %u states", nn_->size());


  return goal->isAchieved();
}

/*
bool TRRT::solve(const base::PlannerTerminationCondition &ptc)
{
  std::cout << "starting solve TRRT" << std::endl;

  checkValidity();
  base::Goal* goal = pdef_->getGoal().get();
  base::GoalSampleableRegion *goal_s = dynamic_cast<base::GoalSampleableRegion*>(goal);

  unsigned int numFail = 0;

  normalizeK_ = 0;

  while(const base::State* st = pis_.nextStart())
  {
    Motion* motion = new Motion(si_);
    si_->copyState(motion->state, st);
    double cost = stateCostEstimator_->computeCost(st);
    normalizeK_ += cost;
    motion->cost = cost;
    nn_->add(motion);

  }

  if(nn_->size() == 0)
  {
    msg_.error("There are no valid initial states!");
    return false;
  }

  normalizeK_ /= nn_->size();

  if(!sampler_)
    sampler_ = si_->allocStateSampler();

  msg_.inform("Starting with %u states", nn_->size());

  Motion *solution  = NULL;
  Motion *approxsol = NULL;
  double  approxdif = std::numeric_limits<double>::infinity();
  Motion *rmotion   = new Motion(si_);
  base::State* rstate = rmotion->state;
  base::State* xstate = si_->allocState();


  while(ptc() == false)
  {
    // sample random state (with goal biasing)
    if (goal_s && rng_.uniform01() < goalBias_ && goal_s->canSample())
        goal_s->sampleGoal(rstate);
    else
        sampler_->sampleUniform(rstate);


    // find closest state in the tree
    Motion* nmotion = nn_->nearest(rmotion);
    base::State* dstate = rstate;

    // find state to add
    double d = si_->distance(nmotion->state, rstate);
    if(d > maxDistance_)
    {
      si_->getStateSpace()->interpolate(nmotion->state, rstate, maxDistance_ / d, xstate);
      dstate = xstate;
      d = maxDistance_;
    }

    if(si_->checkMotion(nmotion->state, dstate))
    {
      // transition test
      double curCost = stateCostEstimator_->computeCost(dstate);
      double parentCost = nmotion->cost;

      bool transitionTestSucceed = false;

      if(curCost <= costMax_)
      {
        if(curCost <= parentCost)
        {
          transitionTestSucceed = true;
        }
        else
        {
          double p = exp(-(curCost - parentCost) / (d * normalizeK_ * temperature_));
          double r = rng_.uniform01();
          if(r < p)
          {
            temperature_ /= temperatureFactor_;
            numFail = 0;
            transitionTestSucceed = true;
          }
          else
          {
            if(numFail > numFailMax_)
            {
              temperature_ *= temperatureFactor_;
              numFail = 0;
            }
            else
              numFail += 1;
          }
        }
      }

      if(transitionTestSucceed)
      {
        Motion *motion = new Motion(si_);
        si_->copyState(motion->state, dstate);
        motion->parent = nmotion;
        motion->cost = curCost;

        nn_->add(motion);
        double dist = 0.0;
        bool solved = goal->isSatisfied(motion->state, &dist);
        if(solved)
        {
          approxdif = dist;
          solution = motion;
          break;
        }

        if(dist < approxdif)
        {
          approxdif = dist;
          approxsol = motion;
        }
      }
    }
  }


  bool approximate = false;
  if(solution == NULL)
  {
    solution = approxsol;
    approximate = true;
  }

  if(solution != NULL)
  {
    // construct the solution path
    std::vector<Motion*> mpath;
    while(solution != NULL)
    {
      mpath.push_back(solution);
      solution = solution->parent;
    }

    // set the solution path
    PathGeometric* path = new PathGeometric(si_);
    for(int i = mpath.size() - 1 ; i >= 0 ; --i)
      path->states.push_back(si_->cloneState(mpath[i]->state));

    goal->setDifference(approxdif);
    goal->setSolutionPath(base::PathPtr(path), approximate);

    if (approximate)
        msg_.warn("Found approximate solution");
  }


  si_->freeState(xstate);
  if(rmotion->state)
    si_->freeState(rmotion->state);
  delete rmotion;

  msg_.inform("Created %u states", nn_->size());

  return goal->isAchieved();
}

*/

}
}
